#include <iostream>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <cassert>
#include <unistd.h>
#include <fcntl.h>
#include <cerrno>
#include <cstring>
#include <string>
#include <thread>
#include <chrono>
#include <atomic>

using namespace std;

const int BUFFER_SIZE = 1024;
atomic<bool> shouldExit{false};  // 线程退出标志

// 设置套接字为非阻塞模式
int setNonBlocking(int fd) {
    int oldOption = fcntl(fd, F_GETFL);
    int newOption = oldOption | O_NONBLOCK;
    fcntl(fd, F_SETFL, newOption);
    return oldOption;
}

// 向服务器发送数据的函数
void sendDataToServer(int sockFd, const string& data) {
    if (send(sockFd, data.c_str(), data.length(), 0) < 0) {
        cerr << "Send data failed: " << strerror(errno) << endl;
    } else {
        cout << "Sent: " << data << endl;
    }
}

// 从服务器接收数据的函数
void receiveDataFromServer(int sockFd) {
    char buffer[BUFFER_SIZE];
    while (!shouldExit) {  // 使用退出标志控制循环
        memset(buffer, 0, BUFFER_SIZE);
        int ret = recv(sockFd, buffer, BUFFER_SIZE - 1, 0);
        if (ret > 0) {
            cout << "Received: " << buffer << endl;
        } else if (ret == 0) {
            cout << "Server closed the connection" << endl;
            shouldExit = true;
            break;
        } else {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                // 非阻塞模式下没有数据可读，短暂休眠
                usleep(100000); // 100ms
                continue;
            }
            cerr << "Receive data failed: " << strerror(errno) << endl;
            shouldExit = true;
            break;
        }
    }
}

int main(int argc, char* argv[]) {
    if (argc <= 2) {
        cout << "Usage: " << argv[0] << " <ip> <port> [--timeout-test]" << endl;
        cout << "  --timeout-test: 测试服务器超时功能" << endl;
        return 1;
    }
    
    const char* ip = argv[1];
    int port = atoi(argv[2]);
    bool isTimeoutTest = (argc > 3 && string(argv[3]) == "--timeout-test");
    
    // 创建套接字
    int sockFd = socket(AF_INET, SOCK_STREAM, 0);
    assert(sockFd >= 0);
    
    // 设置服务器地址
    struct sockaddr_in serverAddress;
    memset(&serverAddress, 0, sizeof(serverAddress));
    serverAddress.sin_family = AF_INET;
    inet_pton(AF_INET, ip, &serverAddress.sin_addr);
    serverAddress.sin_port = htons(port);
    
    // 连接服务器
    int ret = connect(sockFd, (struct sockaddr*)&serverAddress, sizeof(serverAddress));
    if (ret < 0) {
        cerr << "Connect failed: " << strerror(errno) << endl;
        close(sockFd);
        return 1;
    }
    cout << "Connected to server successfully" << endl;
    
    // 设置为非阻塞模式
    setNonBlocking(sockFd);
    
    // 启动接收线程
    thread receiveThread(receiveDataFromServer, sockFd);
    
    if (isTimeoutTest) {
        cout << "Testing timeout feature, will not send data for 20 seconds..." << endl;
        this_thread::sleep_for(chrono::seconds(20));
        cout << "Connection should have been closed by server due to timeout" << endl;
    } else {
        // 正常通信模式
        cout << "Enter 'exit' to quit, or send messages to server:" << endl;
        string message;
        
        // 先发送一条数据，让服务器设置定时器
        sendDataToServer(sockFd, "Hello Server! This is a test message.");
        
        // 交互式发送数据
        while (!shouldExit) {
            getline(cin, message);
            if (message == "exit") {
                shouldExit = true;
                break;
            }
            sendDataToServer(sockFd, message);
        }
        
        // 关闭套接字以通知接收线程
        close(sockFd);
    }
    
    // 等待接收线程结束
    if (receiveThread.joinable()) {
        receiveThread.join();
    }
    
    // 关闭套接字（确保资源释放）
    close(sockFd);
    cout << "Client closed successfully" << endl;
    
    return 0;
}